package com.ztesoft.book.admin.filter;

import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.util.StreamUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;

/**
 * Xss过滤器
 *
 * @author lzq.
 * @description Xss过滤器
 * @date 2021-07-26 17:23
 */
@Slf4j
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    private byte[] requestBody;
    private Charset charSet;

    XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
        // 缓存请求body
        try {
            String requestBodyStr = getRequestPostStr(request);
            if (StrUtil.isNotBlank(requestBodyStr)) {
                if (JSONUtil.isJsonObj(requestBodyStr)) {
                    requestBodyStr = filter(requestBodyStr);
                    JSONObject resultJson = JSONObject.parseObject(requestBodyStr);
                    requestBody = resultJson.toString()
                            .getBytes(charSet);
                }
                else {
                    requestBody = requestBodyStr.getBytes(charSet);
                }
            }
            else {
                requestBody = new byte[0];
            }
        }
        catch (IOException e) {
            log.error("获取body参数异常", e);
        }
    }

    private String filter(String value) {
        if (value != null) {
            value = clearXss(value);
        }
        return value;
    }

    /**
     * 获取body
     *
     * @param request
     *
     * @return
     *
     * @throws IOException
     */
    private String getRequestPostStr(HttpServletRequest request) throws IOException {
        String charSetStr = request.getCharacterEncoding();
        if (charSetStr == null) {
            charSetStr = "UTF-8";
        }
        charSet = Charset.forName(charSetStr);

        return StreamUtils.copyToString(request.getInputStream(), charSet);
    }

    /**
     * 对数组参数进行特殊字符过滤
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values == null) {
            return null;
        }
        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = clearXss(values[i]);
        }
        return encodedValues;
    }

    /**
     * 对参数中特殊字符进行过滤
     */
    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        if (value == null) {
            return null;
        }
        return clearXss(value);
    }

    /**
     * 获取attribute,特殊字符过滤
     */
    @Override
    public Object getAttribute(String name) {
        Object value = super.getAttribute(name);
        if (value instanceof String) {
            clearXss((String) value);
        }
        return value;
    }

    /**
     * 对请求头部进行特殊字符过滤
     */
    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (XssFilterUtil.TOKEN.equals(name)) {
            return clearXss(value);
        }

        if (value == null) {
            return null;
        }
        return clearXss(value);
    }

    /**
     * @param value value
     *
     * @return String
     */
    private String clearXss(String value) {
        if (StringUtils.isEmpty(value)) {
            return value;
        }
        return XssFilterUtil.stripXss(value);
    }

    @Override
    public ServletInputStream getInputStream() {
        if (requestBody == null) {
            requestBody = new byte[0];
        }

        final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(requestBody);

        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener readListener) {
            }

            @Override
            public int read() {
                return byteArrayInputStream.read();
            }
        };
    }

    /**
     * getReader
     *
     * @return BufferedReader
     *
     * @throws IOException IOException
     */
    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream(), StandardCharsets.UTF_8));
    }


}
